"""
LSTM 预测模型
"""

import pandas as pd
import numpy as np
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from math import sqrt
import matplotlib.pyplot as plt
import math

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 读取数据
df = pd.read_excel('data/result.xlsx')
df = df[df['品类'] == '花菜类']
df = df.head(1000)

# 数据归一化
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(df['批发价格(元/千克)'].values.reshape(-1, 1))

# 按7比3划分训练集和测试集
train_size = int(len(scaled_data) * 0.7)
train = scaled_data[:train_size]
test = scaled_data[train_size:]

# 转换数据格式为符合 LSTM 输入要求
def create_dataset(dataset, look_back=30):
    dataX, dataY = [], []
    for i in range(len(dataset) - look_back - 1):
        a = dataset[i:(i + look_back), 0]
        dataX.append(a)
        dataY.append(dataset[i + look_back, 0])
    return np.array(dataX), np.array(dataY)

look_back = 30
trainX, trainY = create_dataset(train, look_back)
testX, testY = create_dataset(test, look_back)

trainX = np.reshape(trainX, (trainX.shape[0], 1, trainX.shape[1]))
testX = np.reshape(testX, (testX.shape[0], 1, testX.shape[1]))

# 构建 LSTM 模型
model = Sequential()
model.add(LSTM(200, return_sequences=True, input_shape=(1, look_back)))
model.add(LSTM(200))
model.add(Dense(1))
model.compile(loss='mean_squared_error', optimizer='adam')
history = model.fit(trainX, trainY, epochs=50, batch_size=1, verbose=2)

# 预测训练集
trainPredict = model.predict(trainX)
trainPredict = scaler.inverse_transform(trainPredict)
trainY_inverse = scaler.inverse_transform([trainY])

# 计算训练集 RMSE
trainScore = math.sqrt(mean_squared_error(trainY_inverse[0], trainPredict[:, 0]))
mse = mean_squared_error(trainY_inverse[0], trainPredict[:, 0])
mae = mean_absolute_error(trainY_inverse[0], trainPredict[:, 0])
r2 = r2_score(trainY_inverse[0], trainPredict[:, 0])

print(f'训练集 MSE: {mse}')
print(f'训练集 MAE: {mae}')
print(f'训练集 R²: {r2}')
print(f'训练集 RMSE: {trainScore:.2f}')

# 绘制训练集预测结果
plt.figure(figsize=(12, 6))
plt.plot(range(look_back, train_size), df['批发价格(元/千克)'].values[look_back:train_size], label='训练集真实值')
plt.plot(range(look_back, len(trainPredict) + look_back), trainPredict, label='训练集预测值', color='green', linestyle='--')

plt.xlabel('时间')
plt.ylabel('批发价格(元/千克)')
plt.legend()
plt.title('训练集预测结果')
plt.grid()
plt.show()

# 预测测试集
testPredict = model.predict(testX)
testPredict = scaler.inverse_transform(testPredict)
testY_inverse = scaler.inverse_transform([testY])

# 计算测试集 RMSE
testScore = math.sqrt(mean_squared_error(testY_inverse[0], testPredict[:, 0]))
mse_test = mean_squared_error(testY_inverse[0], testPredict[:, 0])
mae_test = mean_absolute_error(testY_inverse[0], testPredict[:, 0])
r2_test = r2_score(testY_inverse[0], testPredict[:, 0])

print(f'测试集 MSE: {mse_test}')
print(f'测试集 MAE: {mae_test}')
print(f'测试集 R²: {r2_test}')
print(f'测试集 RMSE: {testScore:.2f}')

# 绘制测试集预测结果
plt.figure(figsize=(12, 6))
plt.plot(range(train_size + look_back, len(df)), df['批发价格(元/千克)'].values[train_size + look_back:], label='测试集真实值')
plt.plot(range(train_size + look_back, train_size + look_back + len(testPredict)), testPredict, label='测试集预测值', color='blue', linestyle='--')

plt.xlabel('时间')
plt.ylabel('批发价格(元/千克)')
plt.legend()
plt.title('测试集预测结果')
plt.grid()
plt.show()

# 预测未来100天的数据
def predict_future(data, model, look_back, days=100):
    predictions = []
    current_input = data[-look_back:]
    for _ in range(days):
        current_input_reshaped = np.reshape(current_input, (1, 1, look_back))
        next_pred = model.predict(current_input_reshaped)
        predictions.append(next_pred[0, 0])
        current_input = np.append(current_input[1:], next_pred)
    return predictions

# 使用测试集最后30个数据预测未来100天
last_30_days = test[-look_back:]
future_predictions = predict_future(last_30_days, model, look_back, days=100)

# 反归一化预测结果
future_predictions = scaler.inverse_transform(np.array(future_predictions).reshape(-1, 1))

# 打印预测结果
print("未来100天的预测结果：", future_predictions.flatten())

# 绘制预测结果
plt.figure(figsize=(12, 6))
plt.plot(df['批发价格(元/千克)'].values, label='真实值')
plt.plot(range(len(df), len(df) + 100), future_predictions, label='预测值', color='red')
plt.xlabel('时间')
plt.ylabel('批发价格(元/千克)')
plt.legend()
plt.title('未来100天的预测结果')
plt.grid()
plt.show()